import heapq
from collections import defaultdict


class NetworkTopology:
    """网络拓扑结构管理"""

    def __init__(self, nodes, links):
        """
        初始化网络拓扑
        :param nodes: 节点列表，如 ['A', 'B', 'C']
        :param links: 链路字典，格式：
            {
                ('A','B'): {'bandwidth': 1e9, 'delay': 0.01},
                ('B','C'): {'bandwidth': 2e9, 'delay': 0.02}
            }
        """
        self.nodes = nodes
        self.original_links = links
        self.reset_links()

    def reset_links(self):
        """重置链路状态"""
        self.available_bandwidth = defaultdict(dict)
        for (u, v), attr in self.original_links.items():
            self.available_bandwidth[u][v] = attr['bandwidth']
            self.available_bandwidth[v][u] = attr['bandwidth']

        self.graph = self._build_graph()

    def _build_graph(self):
        """构建带权图结构"""
        graph = defaultdict(dict)
        for (u, v), attr in self.original_links.items():
            # 权重用剩余带宽的倒数（贪心选择最大带宽路径）
            weight = 1 / (self.available_bandwidth[u][v] + 1e-6)
            graph[u][v] = weight
            graph[v][u] = weight
        return graph


class GreedyRouter:
    """贪心路由算法实现"""

    def __init__(self, topology):
        self.topology = topology
        self.routes = defaultdict(list)

    def find_max_bandwidth_path(self, src, dst):
        """Dijkstra算法寻找最大带宽路径"""
        graph = self.topology.graph

        # 优先队列存储(累计权重, 当前节点, 路径)
        heap = [(0, src, [])]
        visited = set()

        while heap:
            current_cost, current_node, path = heapq.heappop(heap)

            if current_node in visited:
                continue

            path = path + [current_node]

            if current_node == dst:
                return path

            visited.add(current_node)

            for neighbor, weight in graph[current_node].items():
                if neighbor not in visited:
                    heapq.heappush(
                        heap,
                        (current_cost + weight, neighbor, path)
                    )

        raise ValueError(f"No path from {src} to {dst}")

    def allocate_flow(self, flow):
        """分配单个流量"""
        src, dst, demand = flow
        path = self.find_max_bandwidth_path(src, dst)

        # 记录路由路径
        self.routes[(src, dst)].append({
            'path': path,
            'allocated': demand
        })

        # 更新链路带宽
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            self.topology.available_bandwidth[u][v] -= demand
            self.topology.available_bandwidth[v][u] -= demand

        # 更新图结构权重
        self.topology.graph = self.topology._build_graph()

    def calculate_completion_time(self):
        """计算任务完成时间"""
        max_time = 0
        link_usage = defaultdict(float)

        # 统计每条链路的传输时间
        for flow_info in self.routes.values():
            for allocation in flow_info:
                path = allocation['path']
                demand = allocation['allocated']

                for i in range(len(path) - 1):
                    u, v = path[i], path[i + 1]
                    link = (u, v)
                    original_bw = self.topology.original_links[link]['bandwidth']
                    delay = self.topology.original_links[link]['delay']

                    # 传输时间 = 数据量/带宽 + 传播延迟
                    transmission_time = (demand * 8) / original_bw  # 转换为bits
                    total_time = transmission_time + delay

                    link_usage[link] += total_time
                    if link_usage[link] > max_time:
                        max_time = link_usage[link]

        return max_time


def main():
    # 输入示例（可替换为实际数据）
    nodes = ['A', 'B', 'C']
    links = {
        ('A', 'B'): {'bandwidth': 1e9, 'delay': 0.01},  # 1Gbps带宽，10ms延迟
        ('B', 'C'): {'bandwidth': 2e9, 'delay': 0.02},  # 2Gbps带宽，20ms延迟
        ('A', 'C'): {'bandwidth': 1.5e9, 'delay': 0.03}
    }
    traffic_matrix = [
        ('A', 'C', 500e6),  # 500MB流量
        ('B', 'C', 800e6),  # 800MB流量
        ('A', 'B', 300e6)
    ]

    # 初始化拓扑和路由
    topology = NetworkTopology(nodes, links)
    router = GreedyRouter(topology)

    # 按流量矩阵分配路径
    for flow in traffic_matrix:
        router.allocate_flow(flow)

    # 输出结果
    print("路由策略：")
    for (src, dst), allocations in router.routes.items():
        print(f"{src} -> {dst}:")
        for alloc in allocations:
            print(f"  路径: {' → '.join(alloc['path'])}")
            print(f"  分配流量: {alloc['allocated'] / 1e6:.2f} MB")

    total_time = router.calculate_completion_time()
    print(f"\n预计任务完成时间: {total_time:.4f} 秒")


if __name__ == "__main__":
    main()